package com.xx.sparkdemo

import org.apache.spark.sql.catalyst.expressions.Literal
import org.apache.spark.sql.{Column, Dataset, Row, SparkSession}


/**
 *
 * @author tzp
 * @since 2021/7/30
 */
object WithColumnTest {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .master("local")
      .getOrCreate()
    val sc = spark.sparkContext


    import org.apache.spark.sql.functions._
    import spark.implicits._
    val data =
      """
        |sku,store,inhand,storeQuantity
        |1,2222,3,34
        |1,3333,5,45
        |2,4444,5,56
        |2,5555,6,67
        |2,6666,7,67
        |""".stripMargin.lines.toList
    val ds = sc.parallelize(data).toDS()
    val df: Dataset[Row] = spark.read.option("header", true).option("inferSchema", true).csv(ds)

    df.withColumn(
      "storeInfo", struct($"store", struct($"inhand", $"storeQuantity"))).
      groupBy("sku").agg(collect_list("storeInfo").as("info")).
      toJSON.show(false)
    //+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
    //|value                                                                                                                                                                               |
    //+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
    //|{"sku":1,"info":[{"store":2222,"col2":{"inhand":3,"storeQuantity":34}},{"store":3333,"col2":{"inhand":5,"storeQuantity":45}}]}                                                      |
    //|{"sku":2,"info":[{"store":4444,"col2":{"inhand":5,"storeQuantity":56}},{"store":5555,"col2":{"inhand":6,"storeQuantity":67}},{"store":6666,"col2":{"inhand":7,"storeQuantity":67}}]}|
    //+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+



    df.withColumn(
      "storeInfo", struct($"store", struct($"inhand", $"storeQuantity").as("fixedFieldName"))).
      groupBy("sku").agg(collect_list("storeInfo").as("info")).
      toJSON.show(false)
    //+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
    //|value                                                                                                                                                                                              |
    //+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
    //|{"sku":1,"info":[{"store":2222,"myColName":{"inhand":3,"storeQuantity":34}},{"store":3333,"myColName":{"inhand":5,"storeQuantity":45}}]}                                                           |
    //|{"sku":2,"info":[{"store":4444,"myColName":{"inhand":5,"storeQuantity":56}},{"store":5555,"myColName":{"inhand":6,"storeQuantity":67}},{"store":6666,"myColName":{"inhand":7,"storeQuantity":67}}]}|
    //+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

//    df.withColumn(
//      "storeInfo",
//      map($"store", struct($"inhand", $"storeQuantity"), lit("Store"), $"store")).
//      groupBy("sku").agg(collect_list("storeInfo").as("info")).
//    toJSON.show(false)
    // output exception:
    // The given values of function map should all be the same type, but they are [struct<inhand:int,storeQuantity:int>, int]


    df.withColumn(
      "storeInfo",
      map($"store", struct($"inhand", $"storeQuantity"))).
      groupBy("sku").agg(collect_list("storeInfo").as("info")).
      toJSON.show(false)
    //+---------------------------------------------------------------------------------------------------------------------------------------------+
    //|value                                                                                                                                        |
    //+---------------------------------------------------------------------------------------------------------------------------------------------+
    //|{"sku":1,"info":[{"2222":{"inhand":3,"storeQuantity":34}},{"3333":{"inhand":5,"storeQuantity":45}}]}                                         |
    //|{"sku":2,"info":[{"4444":{"inhand":5,"storeQuantity":56}},{"5555":{"inhand":6,"storeQuantity":67}},{"6666":{"inhand":7,"storeQuantity":67}}]}|
    //+---------------------------------------------------------------------------------------------------------------------------------------------+

    df.withColumn("name", lit("fafa")).printSchema()
    val x = Literal.create("fafa")
    x.nullable

//    df.withColumn("name", ).printSchema()

//    df.repartition(200, $"pmod(hash(b)")
//    df.repartition(expr("pmod(hash(b), 200)"))
    df.withColumn("pday", lit("20210718")).write.insertInto("dp_data_db.temp_tzp_bucket5")
  }
}
